{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83c4c5fc-9da0-45b3-9e92-ccafcfcb856f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from Models.OTEModel import Model\n",
    "from Config import config\n",
    "\n",
    "from pytorch_grad_cam import GradCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad\n",
    "from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget\n",
    "from pytorch_grad_cam.utils.image import show_cam_on_image\n",
    "from torchvision.models import resnet50\n",
    "import torchvision\n",
    "import torch\n",
    "from matplotlib import pyplot as plt\n",
    "import numpy as np\n",
    "def myimshows(imgs, titles=False, fname=\"test.jpg\", size=1):\n",
    "    lens = len(imgs)\n",
    "    fig = plt.figure(figsize=(size * lens,size))\n",
    "    if titles == False:\n",
    "        titles=\"0123456789\"\n",
    "    for i in range(1, lens + 1):\n",
    "        cols = 100 + lens * 10 + i\n",
    "        plt.xticks(())\n",
    "        plt.yticks(())\n",
    "        plt.subplot(cols)\n",
    "        if len(imgs[i - 1].shape) == 2:\n",
    "            plt.imshow(imgs[i - 1], cmap='Reds')\n",
    "        else:\n",
    "            plt.imshow(imgs[i - 1])\n",
    "        plt.title(titles[i - 1])\n",
    "    plt.xticks(())\n",
    "    plt.yticks(())\n",
    "    plt.savefig(fname, bbox_inches='tight')\n",
    "    plt.show()\n",
    "def tensor2img(tensor,heatmap=False,shape=(224,224)):\n",
    "    np_arr=tensor.detach().numpy()#[0]\n",
    "    #对数据进行归一化\n",
    "    if np_arr.max()>1 or np_arr.min()<0:\n",
    "        np_arr=np_arr-np_arr.min()\n",
    "        np_arr=np_arr/np_arr.max()\n",
    "    #np_arr=(np_arr*255).astype(np.uint8)\n",
    "    if np_arr.shape[0]==1:\n",
    "        np_arr=np.concatenate([np_arr,np_arr,np_arr],axis=0)\n",
    "    np_arr=np_arr.transpose((1,2,0))\n",
    "    return np_arr\n",
    " \n",
    "path = './1_left.jpg'\n",
    "    \n",
    "bin_data=torchvision.io.read_file(path)#加载二进制数据\n",
    "img=torchvision.io.decode_image(bin_data)/255#解码成CHW的图片\n",
    "img=img.unsqueeze(0)#变成BCHW的数据，B==1; squeeze\n",
    "input_tensor=torchvision.transforms.functional.resize(img,[224, 224])\n",
    " \n",
    "#对图像进行水平翻转，得到两个数据\n",
    "input_tensors=torch.cat([input_tensor, input_tensor.flip(dims=(3,))],axis=0)\n",
    "  \n",
    "model = Model(config).cuda()\n",
    "load_model_path = './save_models/CMAT/pytorch_model.bin' # WAVM\n",
    "model.load_state_dict(torch.load(load_model_path))\n",
    "cv_model = model.img_model\n",
    "\n",
    "target_layers = [cv_model.mamba_model.layers[-1][-2][1].sk_conv]# WAVM\n",
    "\n",
    "with GradCAM(model=cv_model, target_layers=target_layers) as cam:\n",
    "    grayscale_cams = cam(input_tensor=input_tensors)\n",
    "    for grayscale_cam,tensor in zip(grayscale_cams,input_tensors):\n",
    "        #将热力图结果与原图进行融合\n",
    "        rgb_img=tensor2img(tensor)\n",
    "        visualization = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)\n",
    "        myimshows([rgb_img, grayscale_cam, visualization],[\"image\",\"cam\",\"image + cam\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35524ba3-de06-4d7f-b3c7-67f88fae6247",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
